/* 泽天智航 版权所有 */
/* usbmon_capture -- send a USB port reset to a USB device */

/**
 * @file usbmon_capture.c
 * @brief USB设备监控和重置工具
 * 
 * 该程序用于监控USB设备的数据包传输，检测错误并在需要时自动重置设备。
 * 它可以通过多种方式识别USB设备：总线/设备号、厂商/产品ID、序列号或产品名称。
 * 主要功能包括：
 * 1. 列出系统中所有USB设备
 * 2. 根据指定参数查找USB设备
 * 3. 重置USB设备
 * 4. 捕获并分析USB数据包
 * 5. 在检测到特定错误时自动重置设备
 */


#include <stdint.h>
#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <ctype.h>
#include <limits.h>
#include <dirent.h>
#include <sys/ioctl.h>

#include <linux/usbdevice_fs.h>

/**
 * @struct usbmon_packet
 * @brief USB监控事件头结构（mon_bin_hdr）
 * 
 * 该结构体用于表示从usbmon接口读取的USB数据包信息，
 * 包含了USB请求块(URB)的完整信息，如ID、类型、状态和数据长度等。
 */
struct usbmon_packet {
    uint64_t id;          // URB ID - 唯一标识一个USB请求块
    unsigned char type;   // 类型：'S'提交(Submit), 'C'回调(Callback), 'E'错误(Error)
    unsigned char xfer_type; // 传输类型：0=ISO(等时传输),1=INT(中断传输),2=BULK(批量传输),3=CTRL(控制传输)
    unsigned char epnum;  // 端点号 - 标识设备上的通信端点
    unsigned char dev;    // 设备号 - 在总线上的设备地址
    unsigned short busnum; // 总线号 - USB总线编号
    char flag_setup;      // 设置标志 - 控制传输的设置阶段标志
    char flag_data;       // 数据标志 - 指示数据的方向和状态
    int64_t ts_sec;       // 时间戳秒 - 事件发生的秒数部分
    int32_t ts_usec;      // 时间戳微秒 - 事件发生的微秒部分
    int status;           // 状态 - URB完成状态码，0表示成功，负值表示错误
    unsigned int length;  // 数据长度 - 请求/响应的总数据长度
    unsigned int len_cap; // 捕获长度 - 实际捕获的数据长度
    // 对于控制传输，还有8字节设置数据；对于ISO，还有额外字段。这里简化处理
};

/**
 * @struct usbentry
 * @brief USB设备条目结构
 * 
 * 该结构体用于存储USB设备的标识信息和描述信息，
 * 包括总线号、设备号、厂商ID、产品ID以及各种描述性文本。
 */
struct usbentry {
	int bus_num;          // USB总线号
	int dev_num;          // 设备号
	int vendor_id;        // 厂商ID (VID)
	int product_id;       // 产品ID (PID)
	char vendor_name[128]; // 厂商名称
	char product_name[128]; // 产品名称
	char serial[128];     // 序列号
};

/**
 * @brief 从sysfs文件系统读取USB设备属性
 * 
 * 该函数通过sysfs文件系统读取指定USB设备的特定属性值。
 * sysfs是Linux内核提供的一个虚拟文件系统，用于将内核数据结构、设备和驱动程序信息导出到用户空间。
 * 
 * @param dev 设备路径名称，相对于/sys/bus/usb/devices/
 * @param attr 要读取的属性名称
 * @return 成功时返回属性值字符串，失败时返回NULL
 */
static char *sysfs_attr(const char *dev, const char *attr)
{
	int fd, len = 0;
	char path[PATH_MAX];
	static char buf[129];  // 静态缓冲区用于存储读取的属性值

	memset(buf, 0, sizeof(buf));  // 清空缓冲区
	snprintf(path, sizeof(path) - 1, "/sys/bus/usb/devices/%s/%s", dev, attr);  // 构建完整的属性文件路径

	fd = open(path, O_RDONLY);  // 以只读方式打开属性文件
	if (fd >= 0) {  // 如果文件成功打开
		len = read(fd, buf, sizeof(buf) - 1);  // 读取属性值
		close(fd);  // 关闭文件
	}

	while (--len > 0 && isspace(buf[len]))  // 移除尾部的空白字符
		buf[len] = 0;

	return (len >= 0) ? buf : NULL;  // 如果读取成功则返回属性值，否则返回NULL
}

/**
 * @brief 解析USB设备目录并填充设备信息
 * 
 * 该函数从/sys/bus/usb/devices目录中读取USB设备信息，
 * 解析设备的各种属性（如总线号、设备号、厂商ID、产品ID等），
 * 并将这些信息填充到usbentry结构体中。
 * 
 * @param d 已打开的/sys/bus/usb/devices目录的DIR指针
 * @return 成功时返回填充了设备信息的usbentry结构体指针，失败时返回NULL
 * @note 函数使用静态结构体存储设备信息，因此返回值在下次调用前有效
 */
static struct usbentry *parse_devlist(DIR *d)
{
	char *attr;
	struct dirent *e;
	static struct usbentry dev;  // 静态结构体用于存储设备信息

	// 查找有效的USB设备目录项
	// 跳过不以数字开头的目录和包含冒号的目录（接口）
	do {
		e = readdir(d);

		if (!e)  // 如果没有更多目录项，返回NULL
			return NULL;
	} while (!isdigit(e->d_name[0]) || strchr(e->d_name, ':'));

	// 初始化设备结构体
	memset(&dev, 0, sizeof(dev));
	dev.vendor_id = -1;
	dev.product_id = -1;

//	printf("<<<----- %s %d\n", __func__, __LINE__);fflush(stdout);
	// 读取总线号
	attr = sysfs_attr(e->d_name, "busnum");
	if (attr)
		dev.bus_num = strtoul(attr, NULL, 10);

	// 读取设备号
	attr = sysfs_attr(e->d_name, "devnum");
	if (attr)
		dev.dev_num = strtoul(attr, NULL, 10);

	// 读取厂商ID（16进制）
	attr = sysfs_attr(e->d_name, "idVendor");
	if (attr)
		dev.vendor_id = strtoul(attr, NULL, 16);

	// 读取产品ID（16进制）
	attr = sysfs_attr(e->d_name, "idProduct");
	if (attr)
		dev.product_id = strtoul(attr, NULL, 16);

	// 读取厂商名称
	attr = sysfs_attr(e->d_name, "manufacturer");
	if (attr)
		strncpy(dev.vendor_name, attr, sizeof(dev.vendor_name) - 1);

	// 读取产品名称
	attr = sysfs_attr(e->d_name, "product");
	if (attr)
		strncpy(dev.product_name, attr, sizeof(dev.product_name) - 1);

	// 读取序列号
	attr = sysfs_attr(e->d_name, "serial");
	if (attr)
		strncpy(dev.serial, attr, sizeof(dev.serial) - 1);

	// 验证设备信息是否完整有效
	if (dev.bus_num && dev.dev_num && dev.vendor_id >= 0 && dev.product_id >= 0)
		return &dev;

	return NULL;  // 如果设备信息不完整，返回NULL
}

/**
 * @brief 列出系统中所有USB设备
 * 
 * 该函数扫描系统中所有USB设备并以格式化方式打印它们的信息，
 * 包括总线号/设备号、厂商ID/产品ID、序列号（如果有）和产品名称。
 * 输出格式经过对齐处理，使所有设备信息整齐排列。
 */
static void list_devices(void)
{
	DIR *devs = opendir("/sys/bus/usb/devices");
	struct usbentry *dev;
	int max_serial_length = 0;

	if (!devs)
		return;

	/* 计算所有设备中最长序列号的长度，用于格式化输出对齐 */
	while ((dev = parse_devlist(devs)) != NULL) {
		int serial_length = strlen(dev->serial);
		if (serial_length > max_serial_length)
			max_serial_length = serial_length;
	}
	closedir(devs);

	// 重新打开目录，准备第二次遍历
	devs = opendir("/sys/bus/usb/devices");
	if (!devs)
		return;

	// 遍历所有设备并打印信息
	while ((dev = parse_devlist(devs)) != NULL) {
		// 打印总线号/设备号和厂商ID/产品ID
		printf("  Number %03d/%03d  ID %04x:%04x  ", dev->bus_num, dev->dev_num, dev->vendor_id,
		       dev->product_id);
		
		// 打印序列号（如果有）或空白占位符
		if (strlen(dev->serial) > 0)
			printf("SN:%-*s", max_serial_length, dev->serial);
		else
			printf("   %*s", max_serial_length, "");
		
		// 打印产品名称
		printf("  %s\n", dev->product_name);
	}

	closedir(devs);
}

/**
 * @brief 根据指定条件查找USB设备
 * 
 * 该函数根据提供的参数（总线号/设备号、厂商ID/产品ID、序列号或产品名称）
 * 在系统中查找匹配的USB设备。只要满足任一条件即认为找到匹配设备。
 * 
 * @param bus 要查找的总线号指针，NULL表示不使用此条件
 * @param dev 要查找的设备号指针，NULL表示不使用此条件
 * @param vid 要查找的厂商ID指针，NULL表示不使用此条件
 * @param pid 要查找的产品ID指针，NULL表示不使用此条件
 * @param serial 要查找的序列号，NULL表示不使用此条件
 * @param product 要查找的产品名称，NULL表示不使用此条件
 * @return 成功找到设备时返回usbentry结构体指针，未找到时返回NULL
 */
static struct usbentry *find_device(int *bus, int *dev, int *vid, int *pid, const char *serial, const char *product)
{
	DIR *devs = opendir("/sys/bus/usb/devices");

	struct usbentry *e, *match = NULL;

	if (!devs)
		return NULL;

	// 遍历所有USB设备，查找匹配条件的设备
	while ((e = parse_devlist(devs)) != NULL)
		// 任一条件匹配即可：总线号和设备号、厂商ID和产品ID、序列号或产品名称
		if ((bus && (e->bus_num == *bus) && (e->dev_num == *dev)) ||
		    (vid && (e->vendor_id == *vid) && (e->product_id == *pid)) ||
		    (serial && !strcasecmp(e->serial, serial)) || (product && !strcasecmp(e->product_name, product))) {
			match = e;
			break;
		}

	closedir(devs);

	return match;
}

/**
 * @brief 在指定总线上查找USB设备
 * 
 * 该函数与find_device类似，但只在指定的总线上查找设备。
 * 首先确认设备在指定总线上，然后再根据其他条件（厂商ID/产品ID、序列号或产品名称）
 * 进行匹配。这在需要区分相同设备连接到不同USB总线的情况下非常有用。
 * 
 * @param bus 要查找的总线号指针，NULL表示不使用此条件
 * @param dev 要查找的设备号指针（未使用），保留参数以保持与find_device函数接口一致
 * @param vid 要查找的厂商ID指针，NULL表示不使用此条件
 * @param pid 要查找的产品ID指针，NULL表示不使用此条件
 * @param serial 要查找的序列号，NULL表示不使用此条件
 * @param product 要查找的产品名称，NULL表示不使用此条件
 * @return 成功找到设备时返回usbentry结构体指针，未找到时返回NULL
 */
static struct usbentry *find_device_on_bus(int *bus, int *dev, int *vid, int *pid, const char *serial, const char *product)
{
	DIR *devs = opendir("/sys/bus/usb/devices");

	struct usbentry *e, *match = NULL;

	if (!devs)
		return NULL;

	// 遍历所有USB设备
	while ((e = parse_devlist(devs)) != NULL) {
		// 首先确认设备在指定总线上
		if (bus && (e->bus_num == *bus)) {
			// 注释掉的代码显示了之前的匹配逻辑，包括设备号匹配
//			if ((bus && (e->bus_num == *bus) && (e->dev_num == *dev)) ||
//				(vid && (e->vendor_id == *vid) && (e->product_id == *pid)) ||
//				(serial && !strcasecmp(e->serial, serial)) || (product && !strcasecmp(e->product_name, product))) {
			
			// 当前逻辑：在指定总线上，根据厂商ID/产品ID、序列号或产品名称匹配
			if ((vid && (e->vendor_id == *vid) && (e->product_id == *pid)) ||
				(serial && !strcasecmp(e->serial, serial)) || (product && !strcasecmp(e->product_name, product))) {
				match = e;
				break;
			}
		}
	}

	closedir(devs);

	return match;
}

/**
 * @brief 重置指定的USB设备
 * 
 * 该函数通过向USB设备发送USBDEVFS_RESET ioctl命令来重置设备。
 * 重置操作相当于物理断开并重新连接设备，可用于解决设备通信问题或恢复设备到初始状态。
 * 函数会打印重置过程的状态信息。
 * 
 * @param dev 指向要重置的USB设备信息结构体的指针
 */
static void reset_device(struct usbentry *dev)
{
	int fd;
	char path[PATH_MAX];

	// 构建USB设备节点路径
	snprintf(path, sizeof(path) - 1, "/dev/bus/usb/%03d/%03d", dev->bus_num, dev->dev_num);

	printf("Resetting %s ... ", dev->product_name);

	// 以写方式打开USB设备
	fd = open(path, O_WRONLY);
	if (fd > -1) {
		// 发送重置命令
		if (ioctl(fd, USBDEVFS_RESET, 0) < 0)
			printf("failed [%s]\n", strerror(errno));
		else
			printf("ok\n");

		close(fd);
	} else {
		printf("can't open [%s]\n", strerror(errno));
	}
}

/**
 * @brief 捕获USB设备数据包并在检测到错误时自动重置设备
 * 
 * 该函数通过usbmon接口监控指定USB设备的数据包传输，
 * 实时显示数据包信息，并在检测到特定错误（状态码-71）时自动重置设备。
 * 如果设备节点不存在，函数会尝试重新扫描并查找设备。
 * 
 * @param usbdev 指向要监控的USB设备信息结构体的指针
 */
static void usb_capture_reset_device(struct usbentry *usbdev) {
    int bus = usbdev->bus_num;
    int dev = usbdev->dev_num;

    // 打开 /dev/usbmonX，其中X是bus号
    char mon_path[32];
    snprintf(mon_path, sizeof(mon_path), "/dev/usbmon%d", bus);
    int fd = open(mon_path, O_RDONLY);
    if (fd < 0) {
        perror("无法打开 usbmon 接口");
        exit(1);
    }

    printf("开始捕获 Bus %03d Device %03d 的USB数据包... (按 Ctrl+C 退出)\n", bus, dev);

    unsigned long int error_cnt = 0;  // 错误计数器
    // 缓冲区，用于读取事件（最大4KB，足够大多数URB）
    char buffer[4096];
    while (1) {
        // 从usbmon接口读取数据包
        ssize_t len = read(fd, buffer, sizeof(buffer));
        if (len < 0) {
            if (errno == EINTR) continue;  // 被信号中断，继续读取
            perror("读取 usbmon 失败");
            break;
        }
        if (len < sizeof(struct usbmon_packet)) continue;  // 数据包太小，跳过

        // 解析数据包头
        struct usbmon_packet *hdr = (struct usbmon_packet *)buffer;
        // 注释掉的代码是之前的过滤逻辑：if (hdr->dev != dev || hdr->busnum != bus) continue;
        if (hdr->busnum != bus) continue; // 因为过滤后设备号可能会变化，所以只过滤总线号

# if 1
        // 打印数据包基本信息
        printf("\ncur dev ---- Number %03d/%03d  ID %04x:%04x", usbdev->bus_num, usbdev->dev_num, usbdev->vendor_id,
                usbdev->product_id);
        printf("--- USB 数据包 (error: %ld)---\n", error_cnt);
        printf("URB ID: %llu\n", (unsigned long long)hdr->id);
        printf("类型: %c (S=提交, C=回调, E=错误)\n", hdr->type);
        printf("传输类型: %d (0=ISO,1=INT,2=BULK,3=CTRL)\n", hdr->xfer_type);
        printf("端点: 0x%02x\n", hdr->epnum);
        printf("设备: %d (总线: %d)\n", hdr->dev, hdr->busnum);
        printf("时间戳: %ld.%06d\n", (long)hdr->ts_sec, hdr->ts_usec);
        printf("状态: %d\n", hdr->status);
        printf("数据长度: %u (捕获: %u)\n", hdr->length, hdr->len_cap);
#endif

        // 检测错误状态码-71（EPROTO：协议错误）
        if (hdr->status == -71) {
            // char cmd[256] = {};
            // 打印设备复位提示
            printf("设备复位 <<<<----------\n");
            printf("设备复位 <<<<----------\n");
            printf("设备复位 <<<<----------\n");
            printf("设备复位 <<<<----------\n");
            printf("设备复位 <<<<----------\n");
            printf("设备复位 <<<<----------\n");
            // close(fd);  // 注释掉的代码
            error_cnt++;  // 增加错误计数

            // 判断设备节点是否存在
            // 如果不存在则重新扫描（可能是设备被拔掉）
            char path[PATH_MAX] = {};
            snprintf(path, sizeof(path) - 1, "/dev/bus/usb/%03d/%03d", usbdev->bus_num, usbdev->dev_num);
            printf("old----  Number %03d/%03d  ID %04x:%04x  %s\n", usbdev->bus_num, usbdev->dev_num, usbdev->vendor_id,
                    usbdev->product_id, path);fflush(stdout);
            
            // 检查设备节点是否存在
            if (access(path, F_OK) == -1) {
                // 在指定总线上重新扫描设备
                int vid = usbdev->vendor_id;
                int pid = usbdev->product_id;
                printf("重新扫描设备\n");fflush(stdout);
                usbdev = find_device_on_bus(&bus, NULL, &vid, &pid, NULL, NULL);
                if (usbdev) {
                    // 找到设备，打印设备信息
                    printf("  Number %03d/%03d  ID %04x:%04x  ", usbdev->bus_num, usbdev->dev_num, usbdev->vendor_id,
                        usbdev->product_id);fflush(stdout);
                    printf("  %s\n", usbdev->product_name);fflush(stdout);
                } else {
                    // 未找到设备，继续监控
                    printf("未找到指定设备\n");fflush(stdout);
                    continue;
                }
            }

            // 重置设备
            reset_device(usbdev);

            // 注释掉的代码：重新打开usbmon接口
            // fd = open(mon_path, O_RDONLY);
            // if (fd < 0) {
            //     perror("无法打开 usbmon 接口");
            //     exit(1);
            // }
            continue;
        }

        // 解析控制传输的设置数据
        if (hdr->xfer_type == 3 && hdr->flag_setup == 0) {  // 控制传输，有设置数据
            unsigned char *setup = (unsigned char *)(buffer + sizeof(struct usbmon_packet));
            printf("控制设置: bmRequestType=0x%02x, bRequest=0x%02x, wValue=0x%04x, wIndex=0x%04x, wLength=0x%04x\n",
                   setup[0], setup[1], (setup[3] << 8) | setup[2], (setup[5] << 8) | setup[4], (setup[7] << 8) | setup[6]);
            // 调整数据起始位置（跳过8字节的设置数据）
            setup += 8;
            len -= 8;
        } else {
            // 普通数据起始位置
            // unsigned char *setup = (unsigned char *)(buffer + sizeof(struct usbmon_packet));
        }

        // 打印数据内容（如果有）
        if (hdr->len_cap > 0) {
            printf("数据: ");
            unsigned char *data = (unsigned char *)(buffer + sizeof(struct usbmon_packet));
            for (unsigned int i = 0; i < hdr->len_cap && i < 32; i++) {  // 只打印前32字节
                printf("%02x ", data[i]);
            }
            printf("\n");
        }
    }
    close(fd);  // 关闭usbmon接口
}

/**
 * @brief 程序入口函数
 * 
 * 该函数解析命令行参数，根据不同的参数格式识别用户指定的USB设备，
 * 然后重置设备并开始捕获USB数据包。支持四种设备指定方式：
 * 1. 总线号/设备号 (BBB/DDD)
 * 2. 厂商ID/产品ID (VVVV:PPPP)
 * 3. 序列号 (SN:SERIAL)
 * 4. 产品名称 ("Product")
 * 
 * @param argc 命令行参数数量
 * @param argv 命令行参数数组
 * @return 成功返回0，失败返回非0值
 */
int main(int argc, char **argv)
{
	int id1, id2;
	struct usbentry *dev;

	// 解析命令行参数，支持四种格式
	if ((argc == 2) && (sscanf(argv[1], "%3d/%3d", &id1, &id2) == 2)) {
		// 格式1: BBB/DDD - 通过总线号和设备号查找
		dev = find_device(&id1, &id2, NULL, NULL, NULL, NULL);
	} else if ((argc == 2) && (sscanf(argv[1], "%4x:%4x", &id1, &id2) == 2)) {
		// 格式2: VVVV:PPPP - 通过厂商ID和产品ID查找
		dev = find_device(NULL, NULL, &id1, &id2, NULL, NULL);
	} else if ((argc == 2) && (strncmp(argv[1], "SN:", 3) == 0)) {

		// 格式3: SN:SERIAL - 通过序列号查找
		dev = find_device(NULL, NULL, NULL, NULL, &(argv[1])[3], NULL);
	} else if ((argc == 2) && strlen(argv[1]) < 128) {
		// 格式4: "Product" - 通过产品名称查找
		dev = find_device(NULL, NULL, NULL, NULL, NULL, argv[1]);
	} else {
		// 参数格式错误，显示使用说明并列出所有设备
		printf("Usage:\n"
		       "  %s VVVV:PPPP - capture by vendor and product id\n"
		       "  %s BBB/DDD   - capture by bus and device number\n"
		       "  %s SN:SERIAL - capture by serial number\n"
		       "  %s \"Product\" - capture by product name\n\n"
		       "Devices:\n", argv[0], argv[0], argv[0], argv[0]);
		list_devices();
		return 1;
	}

	// 加载usbmon内核模块，用于USB数据包捕获
	int ret = system("modprobe usbmon");
	if (ret != 0) {
		perror("add usbmon error\n");
		return 1;
	}

	// 检查是否找到设备
	if (!dev) {
		fprintf(stderr, "No such device found\n");
		return 1;
	}

	// 重置设备
	reset_device(dev);

	// 开始捕获设备数据包并在错误时自动重置
    usb_capture_reset_device(dev);
	return 0;
}